#!/usr/bin/env python3
'''
todo 结合书，仔细了解损失值的计算方式
'''
from lib import wrappers
from lib import dqn_model

import argparse
import time
import numpy as np
import collections

import torch
import torch.nn as nn
import torch.optim as optim

from tensorboardX import SummaryWriter

# 游戏默认的运行环境
DEFAULT_ENV_NAME = "PongNoFrameskip-v4"
# 游戏训练结束的标志
MEAN_REWARD_BOUND = 19.5

# 伽马值，作用类似于Q值、Bellman公司中的奖励损失
GAMMA = 0.99
# 从样本缓冲区中一次性采样的数量
BATCH_SIZE = 32
# 样本缓冲区的最大容量，使用的deque，当打到最大容量时，会自动删除最早的数据
REPLAY_SIZE = 10000
# 学习率
LEARNING_RATE = 1e-4
# 将训练模型同步到目标模型的频率， 作用有点类似于01代码代码中，更新q值时与旧值之间的权重比例
SYNC_TARGET_FRAMES = 1000
REPLAY_START_SIZE = 10000

# 这里的三个变量共同决定在训练的哪个阶段开始，选择策略中随机策略与自行推理策略之间的权重比例
EPSILON_DECAY_LAST_FRAME = 10**5
EPSILON_START = 1.0
EPSILON_FINAL = 0.02

# 创建一个命名元组，保存从环境中获取的最后一次转换数据
Experience = collections.namedtuple('Experience', field_names=['state', 'action', 'reward', 'done', 'new_state'])

# 经验重放缓冲区，保存每次收集的经验、反馈、动作、场景
# 每次从经验重放缓冲区提取的样本都是随机范围
class ExperienceBuffer:
    def __init__(self, capacity):
        # 缓冲区
        self.buffer = collections.deque(maxlen=capacity)

    def __len__(self):
        # 缓冲区长度
        return len(self.buffer)

    def append(self, experience):
        # 添加经验到缓冲区
        self.buffer.append(experience)

    def sample(self, batch_size):
        '''

        return: 从缓冲区返回的收集的样本，返回的格式：[场景状态数组], [每个场景对应执行的动作数组], [每个场景执行完动作后获得的激励数组], [每个场景执行动作后是否结束标识数组], [每个场景执行完动作后到达的新状态数组]
        '''

        # 从经验缓冲区随机采样batch_size个样本
        indices = np.random.choice(len(self.buffer), batch_size, replace=False)
        states, actions, rewards, dones, next_states = zip(*[self.buffer[idx] for idx in indices])
        return np.array(states), np.array(actions), np.array(rewards, dtype=np.float32), \
               np.array(dones, dtype=np.uint8), np.array(next_states)

# 环境交互代理类
# 主要的作用是，与环境进行交互，并将交互的结果保存到经验重放缓冲区
class Agent:
    def __init__(self, env, exp_buffer):
        '''
        gym 环境交互类

        exp_buffer 经验重放缓冲区
        '''

        self.env = env
        self.exp_buffer = exp_buffer
        self._reset()

    def _reset(self):
        '''
        重置环境
        '''

        self.state = env.reset()
        # 总共获取的激励，重置为0
        self.total_reward = 0.0

    def play_step(self, net, epsilon=0.0, device="cpu"):
        '''
        执行一次交互

        return: 如果游戏结束，则反馈获取的总激励，如果游戏未结束，则反馈None
        '''

        done_reward = None

        if np.random.random() < epsilon:
            # 如果随机数小于epsilon，则执行一次随机动作
            action = env.action_space.sample()
        else:
            # 如果随机数大于epsilon，则由神经网路确定执行的动作
            state_a = np.array([self.state], copy=False)
            state_v = torch.tensor(state_a).to(device)
            q_vals_v = net(state_v)
            _, act_v = torch.max(q_vals_v, dim=1)
            action = int(act_v.item())

        # do step in the environment
        # 执行一次动作，获取一次反馈
        new_state, reward, is_done, _ = self.env.step(action)
        self.total_reward += reward

        # 将获取的反馈保存到经验重放缓冲区
        exp = Experience(self.state, action, reward, is_done, new_state)
        self.exp_buffer.append(exp)
        self.state = new_state
        if is_done:
            # 如果游戏结束，则重置游戏
            done_reward = self.total_reward
            self._reset()
        return done_reward


def calc_loss(batch, net, tgt_net, device="cpu"):
    '''
    计算损失

    batch: 从经验重放缓冲区获取的样本
    net：推理神经网路，主要用于计算梯度
    tat_net: 计算下一个状态的值网络，不参与梯度计算，网络结构和net一致
    device: 执行推理神经网络的设备
    '''

    states, actions, rewards, dones, next_states = batch

    # 将样本转移到指定的设备
    states_v = torch.tensor(states).to(device)
    next_states_v = torch.tensor(next_states).to(device)
    actions_v = torch.tensor(actions, dtype=torch.int64).to(device)
    rewards_v = torch.tensor(rewards).to(device)
    done_mask = torch.ByteTensor(dones).to(device)

    # actions_v 执行的动作值 值范围在[0, 5]之间的数组
    # unsqueeze 增加维度，比如[1, 2]则执行完毕后会变为[[1], [2]]
    # gather net神经网路最后反馈的动作维度为6，表示有6个动作每个动作执行后会得到的Q值，而actions_v表示从网络推理得到的6个动作Q值中选择对应的动作Q值
    # squeeze  选择完各自被推理出来的动作概率后，将最后一个维度去除，比如[[1], [2]]变为[1, 2]
    state_action_values = net(states_v).gather(1, actions_v.unsqueeze(-1)).squeeze(-1)
    # 传入下一个状态后，获取最大的概率（max返回的是列表中最大的值与其所对应的索引）
    next_state_values = tgt_net(next_states_v).max(1)[0]
    # 将已经结束的样本获取的激励设置为0
    # 使用done_mask数组，将指定位置置为0，如果done_mask[n] n位置为0则置不设置为0.0，其他设置为0.0
    # 这步很关键，因为如果不这么做，将导致网路收敛，因为对于已经结束的网路，其是没有下一个状态的q值
    next_state_values[done_mask] = 0.0
    # 使用detach使得该推理值与推理神经网路脱离，不会被pytorch自动计算
    # 因为这个状态值仅用于计算bellman方程中的q值索引，而计算损失时仅需要当前状态即可，而不需要下一个状态
    next_state_values = next_state_values.detach()

    # 弄清楚bellman方程 P75 + P84, 通过下一个状态的最大Q值来更新当前状态的Q值，以此来实现q值的更新
    # 网络推理出来的值，网络推理出来的值是每个动作执行的Q值大小
    # 弄清楚目标网络的作用，目标网络是用来得到下一个状态的Q值，然后根据P84的公式，使下一个状态的最大Q值与net网络的当前状态的Q值相同，得到损失表达式更新网络
    # bellman 方程计算近似值与均方差损失
    expected_state_action_values = next_state_values * GAMMA + rewards_v
    return nn.MSELoss()(state_action_values, expected_state_action_values)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("--env", default=DEFAULT_ENV_NAME,
                        help="Name of the environment, default=" + DEFAULT_ENV_NAME)
    parser.add_argument("--reward", type=float, default=MEAN_REWARD_BOUND,
                        help="Mean reward boundary for stop of training, default=%.2f" % MEAN_REWARD_BOUND)
    args = parser.parse_args()
    device = torch.device("cuda" if args.cuda else "cpu")

    env = wrappers.make_env(args.env)

    # net网络，通过神经网络能够直接得到每个动作执行的Q值，省去上一章中查找计算的
    # net网络的入口参数是当前状态
    net = dqn_model.DQN(env.observation_space.shape, env.action_space.n).to(device)
    # tag_net 网络，和net网络结构一致，入口参数是下一个状态
    # 根据tag_net，可以实现P84的公式，即通过tag_net获取下一个状态的最大Q值，然后根据公式
    # tag_net得到的最大Q值在通过计算式后其结果,应与net网络中根据当前状态和执行的动作得到的Q值相同步到目标
    # 后面的loss损失就是通过这样计算损失，通过不断逼近两者的值来训练网络
    # 为什么这样可以训练网络?
    # 以下是GPT4的回答：
    # 这段代码实现了一个基于深度Q网络（DQN）的智能体，用于解决PongNoFrameskip-v4环境问题。DQN使用两个神经网络：net（在线网络）和tgt_net（目标网络）。在线网络负责学习和更新权重，目标网络用于计算目标Q值。这种设置有助于提高训练稳定性。
    #
    # net和tgt_net的主要区别在于更新频率。在线网络在每个训练步骤中都会更新，而目标网络每隔一定数量的帧（由SYNC_TARGET_FRAMES定义）才会从在线网络更新。这种方式可以减小目标值的波动，使学习过程更稳定。
    #
    # 以下是代码的主要部分解释：
    #
    # 定义一些环境、智能体和训练参数。
    # 定义Experience namedtuple，用于表示智能体与环境交互的单个经验。
    # 创建ExperienceBuffer类，用于存储和采样智能体的经验。
    # 创建Agent类，它是智能体的主要类。智能体在环境中采取行动，并将其经验存储到经验缓冲区中。
    # 实现calc_loss函数，计算在线网络和目标网络之间的损失。
    # 在主函数中，设置环境、网络、优化器等。
    # 在主循环中，智能体与环境交互并更新网络权重。每隔SYNC_TARGET_FRAMES帧，将在线网络的权重复制到目标网络。
    # 如果智能体在最近的100个游戏中达到了预定义的平均奖励阈值，就停止训练。
    # 这种训练方法可以收敛，因为DQN利用了经验回放（Experience Replay）和固定目标网络（Fixed Target Network）两种技术。
    #
    # 经验回放：智能体将其经验存储在经验缓冲区中，并从中随机采样一批数据进行学习。这种方法有助于打破数据之间的关联性，使学习过程更稳定。
    # 固定目标网络：在线网络和目标网络之间的分离使得学习过程更稳定。通过在训练中使用目标网络来计算目标Q值，可以减小目标值的波动，从而提高训练稳定性。在一定数量的帧之后，目标网络才从在线网络更新。这样可以确保目标网络在一段时间内保持不变，有助于稳定训练过程。
    tgt_net = dqn_model.DQN(env.observation_space.shape, env.action_space.n).to(device)
    writer = SummaryWriter(comment="-" + args.env)
    print(net)

    buffer = ExperienceBuffer(REPLAY_SIZE)
    agent = Agent(env, buffer)
    epsilon = EPSILON_START

    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)
    total_rewards = []
    frame_idx = 0
    ts_frame = 0
    ts = time.time()
    best_mean_reward = None

    while True:
        frame_idx += 1
        epsilon = max(EPSILON_FINAL, EPSILON_START - frame_idx / EPSILON_DECAY_LAST_FRAME)

        reward = agent.play_step(net, epsilon, device=device)
        if reward is not None:
            total_rewards.append(reward)
            # 计算帧数，每秒多少帧
            speed = (frame_idx - ts_frame) / (time.time() - ts)
            ts_frame = frame_idx
            ts = time.time()
            # 这里之所以是100，仅仅只是要在tensorboard中显示过去100局的平均激励
            # 与实际的网路梯度计算无关
            mean_reward = np.mean(total_rewards[-100:])
            print("%d: done %d games, mean reward %.3f, eps %.2f, speed %.2f f/s" % (
                frame_idx, len(total_rewards), mean_reward, epsilon,
                speed
            ))
            writer.add_scalar("epsilon", epsilon, frame_idx)
            writer.add_scalar("speed", speed, frame_idx)
            writer.add_scalar("reward_100", mean_reward, frame_idx)
            writer.add_scalar("reward", reward, frame_idx)
            # 将训练的模型保存到本地（每当100局的平均奖励达比上次记录的最大值大时，则记录）
            if best_mean_reward is None or best_mean_reward < mean_reward:
                torch.save(net.state_dict(), args.env + "-best.dat")
                if best_mean_reward is not None:
                    print("Best mean reward updated %.3f -> %.3f, model saved" % (best_mean_reward, mean_reward))
                best_mean_reward = mean_reward
            if mean_reward > args.reward:
                print("Solved in %d frames!" % frame_idx)
                break

        if len(buffer) < REPLAY_START_SIZE:
            continue

        # 指定每SYNC_TARGET_FRAMES帧将训练的神经网络同步到目标网路
        if frame_idx % SYNC_TARGET_FRAMES == 0:
            tgt_net.load_state_dict(net.state_dict())

        # 计算net 推理网络的梯度并更新梯度
        optimizer.zero_grad()
        batch = buffer.sample(BATCH_SIZE)
        loss_t = calc_loss(batch, net, tgt_net, device=device)
        loss_t.backward()
        optimizer.step()
    writer.close()
